# -*- coding: utf-8 -*-
"""
Created on Mon Nov 13 11:50:55 2017

@author: za-xuzhaoye
"""

#cointegrated_test

import statsmodels.api as sm
import seaborn as sns
import numpy as np
import pandas as pd



def find_cointegrated_pairs(dataframe):
    n=dataframe.shape[1]
    pvalue_matrix=np.ones((n,n))
    keys=dataframe.keys()
    pairs=[]
    
    for i in range(n):
        for j in range(i+1, n):
            etf1=dataframe[keys[i]]
            etf2=dataframe[keys[j]]
            result=sm.tsa.stattools.coint(etf1,etf2)
            pvalue=result[1]
            pvalue_matrix[i,j]=pvalue
            
            if pvalue<0.05:
                pairs.append((keys[i],keys[j],pvalue))
            
    return pvalue_matrix, pairs
        
#closeprice=open("D:\FOF_strategy\A&H_ETF_data\ClosePrice.csv",'rt')
#price_list=[]
#try:
#    reader=csv.reader(closeprice)
#    for row in reader:
#        price_list.append(row)
#
#except Exception as e:
#    print("Exception is:",e)
#finally:
#    closeprice.close()
#
#for x in range(0,621):
#    etf_list=dict(zip(x))
    
    
data=pd.read_csv('D:\FOF_strategy\A&H_ETF_data\ClosePrice.csv')
#etf=data.set_index（'50ETF','H_ETF','HZETF','SHKETF','150176HGB','159920HSETF','160717HSHG','502050SZ50B','510880HLETF','510880HLETF').
#etf=data[['50ETF','H_ETF','HZETF','SHKETF','150176HGB','159920HSETF','160717HSHG','502050SZ50B','510880HLETF','510880HLETF']]

pvalues,pairs=find_cointegrated_pairs(data)
sns.heatmap(1-pvalues,xticklabels=data,yticklabels=data,cmap='RdYlGn_r',mask=(pvalues==1))


